import { ITensorAllocator } from "../../Core/Backends/ITensorAllocator";
import { Model } from "../../Core/Model";
import { ActionSpec } from "../Actuators/ActionSpec";
import { ActionBuffers } from "../Actuators/IActionReceiver";
import { ContinuousActionOutputApplier, DiscreteActionOutputApplier, LegacyDiscreteActionOutputApplier, MemoryOutputApplier } from "./ApplierImpl";
import { ModelApiVersion } from "./BarracudaModelParamLoader";
import { ModelExtensions } from "./ModelExtensions";
import { TensorNames } from "./TensorNames";
import { TensorProxy } from "./TensorProxy";


export interface IApplier {

    /**
     * @en
     * Applies the values in the Tensor to the Agents present in the agentInfos
     * 
     * @zh
     * 将张量中的值应用于 agentInfos 中存在的代理
     * @param tensorProxy 
     * @param actionIds 
     * @param lastActions 
     */
    apply(tensorProxy: TensorProxy, actionIds: number[], lastActions: { [id: number]: ActionBuffers }): void;
}

/**
 * @en
 * Mapping between the output tensor names and the method that will use the
 * output tensors and the Agents present in the batch to update their action, memories and
 * value estimates.
 * A TensorApplier implements a Dictionary of strings (node names) to an Action.
 * This action takes as input the tensor and the Dictionary of Agent to AgentInfo for
 * the current batch.
 * 
 * @zh
 * 张量灌输器
 * 
 * 输出张量名称和方法之间的映射，该方法将使用输出张量和批次中存在的代理来更新它们的动作、记忆和值估计。
 * 一个 TensorApplier 实现了一个字符串（节点名称）到一个 Action 的字典。
 * 这个 Action 接受 Tensor 和 Agent 到 AgentInfo 字典作为输入，用于当前批次。
 * 
 */
export class TensorApplier {
    private readonly _dict: { [name: string]: IApplier } = {};

    constructor(
        actionSpec: ActionSpec,
        seed: number,
        allocator: ITensorAllocator,
        memories: { [id: number]: number[] },
        barracudaModel: Model | null = null,
        deterministicInference: boolean = false
    ) {
        if (barracudaModel === null) {
            return;
        }

        const model: Model = barracudaModel;
        if (!ModelExtensions.supportsContinuousAndDiscrete(model)) {
            actionSpec.checkAllContinuousOrDiscrete();
        }
        if (actionSpec.numContinuousActions > 0) {
            const tensorName = ModelExtensions.continuousOutputName(model, deterministicInference);
            this._dict[tensorName] = new ContinuousActionOutputApplier(actionSpec);
        }
        const modelVersion = ModelExtensions.getVersion(model);
        if (actionSpec.numDiscreteActions > 0) {
            const tensorName = ModelExtensions.discreteOutputName(model, deterministicInference);
            if (modelVersion === ModelApiVersion.MLAgents1_0) {
                this._dict[tensorName] = new LegacyDiscreteActionOutputApplier(actionSpec, seed, allocator);
            }
            if (modelVersion === ModelApiVersion.MLAgents2_0) {
                this._dict[tensorName] = new DiscreteActionOutputApplier(actionSpec, seed, allocator);
            }
        }
        this._dict[TensorNames.recurrentOutput] = new MemoryOutputApplier(memories);
    }

    applyTensors(tensors: TensorProxy[], actionIds: number[], lastActions: { [id: number]: ActionBuffers }): void {
        for (let tensorIndex = 0; tensorIndex < tensors.length; tensorIndex++) {
            const tensor = tensors[tensorIndex];
            if (!Reflect.has(this._dict, tensor.name)) {
                throw new Error(`Unknown tensorProxy expected as output : ${tensor.name}`);
            }
            this._dict[tensor.name].apply(tensor, actionIds, lastActions);
        }
    }
}